// tasmmeanshape.cpp: Generate mean shape in the face detector frame.
//      This also prints an analysis of face dets in the shapefile.
//      If there are few or no facedet results in the shapefile,
//      it will also generate facedet.shape.
//
//      TODO Would like to simplify this code.
//
// Copyright (C) 2005-2013, Stephen Milborrow

#include "stasm.h"
#include "../tasm/src/tasm.h"

#include <opencv2/highgui/highgui.hpp>

namespace stasm
{
static DetPar VecAsDetPar(
    const VEC& vec,
    EYAW       eyaw)
{
    DetPar detpar;   // detpar constructor sets all fields INVALID
    detpar.eyaw = eyaw;
    if (NSIZE(vec) > 0)  detpar.x      = vec(0);
    if (NSIZE(vec) > 1)  detpar.y      = vec(1);
    if (NSIZE(vec) > 2)  detpar.width  = vec(2);
    if (NSIZE(vec) > 3)  detpar.height = vec(3);
    if (NSIZE(vec) > 4)  detpar.lex    = vec(4);
    if (NSIZE(vec) > 5)  detpar.ley    = vec(5);
    if (NSIZE(vec) > 6)  detpar.rex    = vec(6);
    if (NSIZE(vec) > 7)  detpar.rey    = vec(7);
    if (NSIZE(vec) > 8)  detpar.mouthx = vec(8);
    if (NSIZE(vec) > 9)  detpar.mouthy = vec(9);
    if (NSIZE(vec) > 10) detpar.rot    = vec(10);
    if (NSIZE(vec) > 11) detpar.yaw    = vec(11);
    return detpar;
}

static void WriteDetPar(           // write detpar to facedet.shape
    const DetPar& detpar,          // in
    const char*   base,            // in: from shapefile
    const char*   dirs,            // in: from shapefile
    const char*   facedet_outpath, // in: full pathname of generated facedet.shape file
    const char*   modname,         // in: e.g. "yaw00"
    const char*   cmdline)         // in: command line used to invoke tasm
{
    CV_Assert(detpar.x);
    CV_Assert(facedet_outpath);
    static FILE* file;
    if (!file) // first time?
    {
        lprintf("\nGenerating %s ", facedet_outpath);
        file = fopen(facedet_outpath, "wb");
        if (!file)
            Err("Cannot open %s for writing", facedet_outpath);
        Fprintf(file, "shape %s\n\n", facedet_outpath);
        Fprintf(file, "# Face detector results generated by tasm\n");
        Fprintf(file, "# Command: %s\n", cmdline);
        Fprintf(file, "# Model name: %s\n\n", modname);
        Fprintf(file, "Directories %s\n\n", dirs);
    }
    Fprintf(file, "81000000 %s\n{ 1 10", base);
    Fprintf(file, " %d", cvRound(detpar.x));
    Fprintf(file, " %d", cvRound(detpar.y));
    Fprintf(file, " %d", cvRound(detpar.width));
    Fprintf(file, " %d", cvRound(detpar.height));
    Fprintf(file, " %g", detpar.lex);
    Fprintf(file, " %g", detpar.ley);
    Fprintf(file, " %g", detpar.rex);
    Fprintf(file, " %g", detpar.rey);
    Fprintf(file, " %g", detpar.mouthx);
    Fprintf(file, " %g", detpar.mouthy);
    Fprintf(file, "\n}\n");
    fflush(file);
}

static Shape AlignShapeToDetFrame( // align shape to detpar
    const Shape&  shape,           // in
    const DetPar& detpar)          // in
{
    Shape newshape = shape.clone();

    TransformShapeInPlace(newshape, 1, 0, -detpar.x,
                                    0, 1, -detpar.y);

    const double xscale = DET_FACE_WIDTH / detpar.width;
    const double yscale = DET_FACE_WIDTH / detpar.height;

    TransformShapeInPlace(newshape, xscale, 0,      0,
                                    0,      yscale, 0 );

    return newshape;
}

// If the facedet rectangle is not in the shape file, should we search the
// image for the face?  Use a 60% heuristic to determine this.  We don't
// want to search the image if only a few face rectangles are missing from
// the shape file.  The assumption in that case is that we pre-found the
// facedet rectangles and manually put them the shape file, but the facedet
// failed on a few images, so don't waste time now re-searching those images.

static bool MustSearchImgIfFaceDetNotInShapefile(
    const ShapeFile& sh,                          // in
    const vec_int&   nused_points,                // in
    const char*      modname)                     // in: e.g. "yaw00"
{
    const map_mat* yaws = GetYawMap(sh, modname);
    if (NSIZE(*yaws) == 0) // no facedets in shapefile?
    {
        lprintf("\n    no facedets in the shapefile");
        return true;           // note return
    }
    int nshapes_with_all_points = 0;
    int ndets_shapefile = 0;
    const int npoints = sh.shapes_[0].rows;
    for (int ishape = 0; ishape < sh.nshapes_; ishape++)
    {
        if (nused_points[ishape] == npoints)
        {
            nshapes_with_all_points++;
            const VEC vec(sh.Yaw00_(sh.bases_[ishape])); // get facedet from shapefile
            if (Valid(vec(0)))                           // got it?
                ndets_shapefile++;
        }
    }
    lprintf("\n    %.2f%% of shapes with all landmarks have a facedet in the shapefile ",
             100. * ndets_shapefile / nshapes_with_all_points);

    return double(ndets_shapefile) / nshapes_with_all_points < .60;
}

static void PossiblyDetectEyesAndMouth(
    DetPar&      detpar, // io: eye and mouth possibly updated, other fields untouched
    const Image& img)    // in
{
    static bool firsttime = true;
    if (firsttime)
    {
        firsttime = false;
        if (TASM_DET_EYES_AND_MOUTH)
        {
            if (print_g) // no tasm -q flag?
                printf("\n    opening eye and mouth detectors "
                       "(TASM_DET_EYES_AND_MOUTH is true)");
            OpenEyeMouthDetectors(true /*need_eyes*/, true /*need_mouth*/, TASM_DATADIR);
        }
        else
            if (print_g) // no tasm -q flag?
                printf("\n    will not search for eyes or mouth "
                       "(TASM_DET_EYES_AND_MOUTH is false)");
        logprintf("\n");
    }
    if (TASM_DET_EYES_AND_MOUTH)
        DetectEyesAndMouth(detpar, img);
}

static void GetTasmDetParFromImage(
    DetPar&          detpar,          // out: face detector rectangle and eye and mouth positions
    int&             ndets_image,     // io
    int&             nfalse_pos,      // io: count of false positives by the face detector
    const ShapeFile& sh,              // in
    const Shape&     shape,           // in
    const char*      base,            // in
    const char*      facedet_outpath, // in: if not NULL, search actual image (not just shapefile)
    const char*      modname,         // in: e.g. "yaw00"
    const char*      cmdline,         // in: command line used to invoke tasm
    bool             force_facedet,   // in: facedet the images themselves (ignore shapefile facedets)
    bool             write_detpars)   // in: write the image detpars to facedet.shape
{
    // load the image from disk (must first get its full path)
    const char* imgpath = PathGivenDirs(base, sh.dirs_, sh.shapepath_);
    Image img(cv::imread(imgpath, CV_LOAD_IMAGE_GRAYSCALE));
    if (!img.data)
        Err("Cannot load", imgpath);
    // open facedet if not already open
    static FaceDet facedet;
    static bool firsttime = true;
    if (firsttime)
    {
        firsttime = false;
        if (print_g) // no tasm -q flag?
            printf("\n    opening face detector (TASM_FACEDET_MINWIDTH is %d)",
                   TASM_FACEDET_MINWIDTH);
        logprintf("\n");
        facedet.OpenFaceDetector_(TASM_DATADIR, NULL);
    }
    // call the facedet to search the image
    facedet.DetectFaces_(img, imgpath, 0 /*multiface*/, TASM_FACEDET_MINWIDTH, NULL);
    detpar = facedet.NextFace_();
    if (!Valid(detpar.x))                   // face not found in image?
    {
        const map_mat* yaws = GetYawMap(sh, modname);
        static int printed;
        if (!force_facedet && NSIZE(*yaws)) // looked in  shapefile?
            PrintOnce(printed, "\n    %s: no facedet in the shapefile "
                      "or in the image, skipping... ", base);
        else
            PrintOnce(printed, "\n    %s: no facedet in the image, skipping... ", base);
        return;                             // return failure
    }
    ndets_image++;
    static int printed;
    PrintOnce(printed, "\n    %s: found face in the image... ", base);
    if (FaceDetFalsePos(detpar, shape))     // false positive?
    {
        nfalse_pos++;
        static int printed;
        PrintOnce(printed, "\n    %s: false positive in the image, skipping... ", base);
        const DetPar invalid_detpar;        // all fields marked invalid by default
        detpar = invalid_detpar;            // all fields invalid
        return;                             // return failure
    }
    PossiblyDetectEyesAndMouth(detpar, img); // possibly update eye and mouth fields
    if (write_detpars)
        WriteDetPar(detpar, base, sh.dirs_, facedet_outpath, modname, cmdline);
}

// Get face detector parameters, first lokking in the shapefile.  If they
// are not in the shapefile and facedet_outpath is not NULL, then search
// the image itself.
// Also print progress, and update stat counters for printing later.
// TODO Extend the yaw handling to handle other yaws, not just Yaw00.

static void GetTasmDetPar(
    DetPar&          detpar,          // out: face det rect and eye and mouths
    int&             ndets_shapefile, // io
    int&             ndets_image,     // io
    int&             nfalse_pos,      // io: count of false positives by the face detector
    const ShapeFile& sh,              // in
    const Shape&     shape,           // in
    const char*      base,            // in
    const char*      modname,         // in: e.g. "yaw00"
    const char*      cmdline,         // in: command line used to invoke tasm
    const char*      facedet_outpath, // in: if not NULL, search actual image (not just shapefile)
    bool             force_facedet,   // in: facedet the images themselves (ignore shapefile facedets)
    bool             write_detpars)   // in: write the image detpar to facedet.shape
{
    // there are multiple returns to avoid deep nesting
    const DetPar invalid_detpar;            // all fields marked invalid by default
    detpar = invalid_detpar;                // all fields invalid
    if (!force_facedet)
    {
        const map_mat* yaws = GetYawMap(sh, modname);
        if (NSIZE(*yaws))                   // at least one facedet in shapefile?
        {
            const VEC vec(sh.Yaw00_(base)); // get shapefile facedet
            if (!Valid(vec(0)))             // did not get it?
            {
                static int printed;
                PrintOnce(printed, "\n    %s: no facedet in the shapefile, skipping... ", base);
                return;                     // return failure
            }
            ndets_shapefile++;
            detpar = VecAsDetPar(vec, EYAW00);  // convert from vector to DetPar
            if (FaceDetFalsePos(detpar, shape)) // false positive?
            {
                nfalse_pos++;
                static int printed;
                PrintOnce(printed, "\n    %s: false positive in shapefile, skipping... ", base);
                detpar = invalid_detpar;    // all fields invalid
                return;                     // return failure
            }
            return;                         // return success
        }
        // do not yet have the facdet
        if (!facedet_outpath)               // must not search image for facedet?
            return;                         // return failure
    }
    GetTasmDetParFromImage(detpar, ndets_image, nfalse_pos,
                           sh, shape, base,
                           facedet_outpath, modname, cmdline,
                           force_facedet, write_detpars);
}

// these eye and mouth stats are purely informational and do not affect the model

static void UpdateEyeMouthStats(
    int&          nno_eyes,  // io
    int&          nno_leye,  // io
    int&          nno_reye,  // io
    int&          nno_mouth, // io
    const DetPar& detpar,    // in
    const char*   base)      // in
{
    if (!Valid(detpar.lex) && !Valid(detpar.rex))
    {
        nno_eyes++;
        static int printed;
        PrintOnce(printed,
                 "\n    %s: both eyes not detected... ", base);
    }
    else if (!Valid(detpar.lex))
    {
        nno_leye++;
        static int printed;
        PrintOnce(printed,
                 "\n    %s: left eye not detected... ", base);
    }
    else if (!Valid(detpar.rex))
    {
        nno_reye++;
        static int printed;
        PrintOnce(printed,
                 "\n    %s: right eye not detected... ", base);
    }
    if (!Valid(detpar.mouthx))
    {
        nno_mouth++;
        static int printed;
        PrintOnce(printed,
                 "\n    %s: mouth not detected... ", base);
    }
}

static void ShowDetStats(
    int  nshapes,                // in
    int  nshapes_used,           // in
    int  nshapes_missing_points, // in
    int  ndets_shapefile,        // in
    int  ndets_image,            // in
    int  nfalse_pos,             // in
    int  nno_eyes,               // in
    int  nno_leye,               // in
    int  nno_reye,               // in
    int  nno_mouth,              // in
    bool force_facedet,          // in
    bool mustsearchimg)          // in
{
    const int ndets = ndets_shapefile + ndets_image;
    if (ndets == 0)
    {
        lprintf("\n");
        if (force_facedet)
            Err("No faces detected in the images");
        else if (mustsearchimg)
            Err("No facedets in the shapefile or in the images themselves");
        else
            Err("No facedets in the shapefile");
    }
    else
    {
        lprintf("\nDone generating mean shape aligned to the facedets\n");
        lprintf("    %d face%s detected (%.2f%%), ",
                ndets, plural(ndets), 100. * ndets / nshapes);
        lprintf("%d facedet%s in the shapefile, %d facedet%s in the images\n",
                ndets_shapefile, plural(ndets_shapefile), ndets_image,
                plural(ndets_image));
        lprintf("    %d facedet false positive%s, %d shape%s with missing landmarks\n",
                nfalse_pos, plural(nfalse_pos), nshapes_missing_points,
                plural(nshapes_missing_points));
        lprintf("    %d face%s actually used "
                "(valid facedet%s) (%.2f%% of the shapes in the shapefile)",
                nshapes_used, plural(nshapes_used),
                TASM_MEANSHAPE_USES_SHAPES_WITH_MISSING_POINTS? "": " and all points",
                100. * nshapes_used / nshapes);
        if (TASM_DET_EYES_AND_MOUTH)
            lprintf("\n    %d missing both eyes, %d missing just left eye, "
                    "%d missing just right eye, %d missing mouth",
                    nno_eyes, nno_leye, nno_reye, nno_mouth);
    }
    lprintf("\n");
    if (nshapes_used == 0)
    {
        if (!force_facedet && mustsearchimg)
            Err("No valid facedets in the shapefile or in the images themselves");
        else
            Err("No valid facedets in the shapefile");
    }
}

// Generate meanshape which is the mean shape of shapes in the
// face dector frame, aligned to the facedet box.
// This alignment to the face det frame uses an idea by GuoQing Hu.

void MeanShapeAlignedToFaceDets(
    Shape&           meanshape,     // out
    const ShapeFile& sh,            // in
    const vec_int&   nused_points,  // in
    const char*      outdir,        // in: output directory (-d flag)
    const char*      modname,       // in: e.g. "yaw00"
    const char*      cmdline,       // in: command line used to invoke tasm
    bool             force_facedet, // in: facedet the images themselves (ignore shapefile facedets)
    bool             write_detpars) // in: write the image detpars to facedet.shape
{
    lprintf("Generating mean shape aligned to the facedets ");

    // if we need the face detector and to generate facedet.shape, set facedet_outpath
    char path[SBIG];
    char* facedet_outpath = NULL;
    if (force_facedet ||
        MustSearchImgIfFaceDetNotInShapefile(sh, nused_points, modname))
    {
        sprintf(path, "%s/log/%s", outdir, "facedet.shape");
        facedet_outpath = path;
    }
    else
        lprintf("\n    will not run the facedet on the actual images "
                "(will use just the facedets saved in the shapefile) ");

    const int npoints = sh.shapes_[0].rows;
    meanshape = MAT(npoints, 2, 0.); // redimension and set all elements to zero
    int nshapes_used = 0;            // nbr of shapes actually used
    int nshapes_missing_points = 0;  // nbr of shapes with missing points
    int ndets_shapefile = 0;         // nbr of face detects from shape file
    int ndets_image = 0;             // nbr of face detects from image itself
    int nfalse_pos = 0;              // nbr of false positives by the face detector
    int nno_eyes = 0, nno_leye = 0, nno_reye = 0, nno_mouth = 0; // eye and mouth stats, purely info
    Pacifier pacifier(sh.nshapes_);
    for (int ishape = 0; ishape < sh.nshapes_; ishape++)
    {
        const char* base = sh.bases_[ishape].c_str();
        const Shape& shape = sh.shapes_[ishape];
        logprintf("\n%s:", base);

        DetPar detpar; // face detector parameters from shapefile or image itself
        GetTasmDetPar(detpar, ndets_shapefile,  ndets_image, nfalse_pos,
                      sh, shape, base, modname, cmdline,
                      facedet_outpath, force_facedet, write_detpars);

        if (Valid(detpar.x) &&         // facedet valid?
            (TASM_MEANSHAPE_USES_SHAPES_WITH_MISSING_POINTS?
                1: nused_points[ishape] == npoints))
        {
            // it's a good shape, with a valid facedet rectangle
            nshapes_used++;
            meanshape += AlignShapeToDetFrame(shape, detpar);
        }
        if (nused_points[ishape] != npoints) // shape does not have all points?
        {
            nshapes_missing_points++;
            static int printed;
            PrintOnce(printed, "\n    %s: shape does not have all landmarks... ", base);
        }
        if (Valid(detpar.x) && TASM_DET_EYES_AND_MOUTH)
        {
            UpdateEyeMouthStats(nno_eyes, nno_leye, nno_reye, nno_mouth,
                                detpar, base);
        }
        if (print_g && facedet_outpath) // print pacifier if searching the actual images
            pacifier.Print_(ishape);
    }
    if (print_g && facedet_outpath)
        pacifier.End_();

    if (nshapes_used)
        meanshape /= nshapes_used;

    ShowDetStats(sh.nshapes_,
                 nshapes_used, nshapes_missing_points, ndets_shapefile,
                 ndets_image, nfalse_pos,
                 nno_eyes, nno_leye, nno_reye, nno_mouth,
                 force_facedet, facedet_outpath != NULL);
}

} // namespace stasm
